package cn.brucemaa.spring_ai_demo.controller;

import com.alibaba.cloud.ai.dashscope.api.DashScopeApi;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatModel;
import com.alibaba.cloud.ai.dashscope.chat.DashScopeChatOptions;
import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingModel;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.util.MimeTypeUtils;
import org.springframework.web.bind.annotation.*;
import reactor.core.publisher.Flux;

import java.net.MalformedURLException;
import java.net.URL;
import java.util.List;

@Slf4j
@RestController
@RequestMapping(path = "/demo")
public class DemoController {

    private static final String API_KEY = "sk-eb8cd962b7fe4930ad380d74ff24e6bc";
    private static final String DEFAULT_PROMPT_MESSAGE = "你好，我是谁？";

    @GetMapping
    public ResponseEntity<String> home() {
        return ResponseEntity.ok("Hello, Spring AI Demo!");
    }

    @PostMapping(path = "/chat")
    public ResponseEntity<String> chat() throws MalformedURLException {
        URL url = new URL("https://docs.spring.io/spring-ai/reference/_images/multimodal.test.png");
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        final DashScopeChatModel dashScopeChatModel = new DashScopeChatModel(dashScopeApi);
        final ChatClient chatClient = ChatClient.builder(dashScopeChatModel).build();
        return ResponseEntity.ok(
                chatClient
                        .prompt()
                        .options(DashScopeChatOptions.builder().build())
                        .user(u -> u.text("解释下你在图片里看到了什么?").media(MimeTypeUtils.IMAGE_PNG, url))
                        .call()
                        .content()
        );
    }

    @PostMapping(path = "/chat0", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ChatResponse> chat0(@RequestBody String message) {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        final DashScopeChatModel dashScopeChatModel = new DashScopeChatModel(dashScopeApi);
        final ChatClient chatClient = ChatClient.builder(dashScopeChatModel).build();
        return chatClient.prompt(message).stream().chatResponse();
    }

    @PostMapping(path = "/chat1")
    public ResponseEntity<String> chat1() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);

        final DashScopeChatModel dashScopeChatModel = new DashScopeChatModel(dashScopeApi);
        return ResponseEntity.ok(dashScopeChatModel.call(DEFAULT_PROMPT_MESSAGE));
    }

    @PostMapping(path = "/chat2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<String> chat2() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        final DashScopeChatModel dashScopeChatModel = new DashScopeChatModel(dashScopeApi);
        return dashScopeChatModel.stream(DEFAULT_PROMPT_MESSAGE);
    }

    @PostMapping(path = "/embedding")
    public EmbeddingResponse embedding() {
        DashScopeApi dashScopeApi = new DashScopeApi(API_KEY);
        var dashScopeEmbeddingModel = new DashScopeEmbeddingModel(dashScopeApi);
        return dashScopeEmbeddingModel.embedForResponse(List.of("Hello World", "World is big and salvation is near"));
    }
}
